import json
import os
import pandas as pd

import pymysql
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from pydantic import BaseModel, Field
from src.agent.my_llm import llm

# base_url = os.environ.get("DEEPSEEK_BASE_URL")
# model_name = os.environ.get("DEEPSEEK_MODEL_NAME")
# os.environ["OPENAI_API_KEY"] = os.getenv("DEEPSEEK_API_KEY")
# tarvily_api_key = os.environ.get("TARVILY_API_KEY")
#
# llm = ChatOpenAI(base_url=base_url, model_name=model_name)

host = "linux.server"
user = "root"
password = "Chenhao$123456"
db = "telco_db"
port = 3306

df = None


# 定义结构化参数模型
class SQLQuerySchema(BaseModel):
    sql_query: str = Field(
        description="字符串形式的SQL查询语句，用于执行对mysql中telco_db数据库中各张表进行查询，并获得各表中的各类数据")


# 封装LangGraph工具
@tool(args_schema=SQLQuerySchema)
def sql_inter(sql_query: str) -> str:
    """
    当需要对数据库进行操作时，可以使用本工具。本工具可以对数据库表进行增、删、改、查。
    该工具用于在指定mysql服务器上运行一段sql代码，完成数据CRUD相关工作，并且当前工具是使用pymysql连接mysql数据库。
    :param sql_query: 字符串形式的SQL语句，用于执行对mysql中telco_db数据库中各张表进行CRUD
    :return: sql_query在mysql中的运行结果
    """
    print("call sql query")
    connection = pymysql.connect(host=host, user=user, password=password, db=db, port=int(port), charset="utf8")
    try:
        with connection.cursor() as cursor:
            cursor.execute(sql_query)
            connection.commit()
            result = cursor.fetchall()
    finally:
        connection.close()
    return json.dumps(result, ensure_ascii=False)


# 定义结构化参数
class ExtractQuerySchema(BaseModel):
    sql_query: str = Field(description="用于从mysql提取数据的sql查询语句。")


# 注册为agent工具
@tool(args_schema=ExtractQuerySchema)
def extract_data(sql_query: str) -> str:
    """
    用于在mysql数据库中提取一张表中的数据到一个全局变量中，这个全局变量是一个pandas data frame类型
    注意：本工具只负责提取数据表，不负责数据查询。若需要再mysql中进行数据查询，请使用sql_inter工具。
    注意：编写外部工具的参数消息时，必须是满足json格式的字符串。
    :param sql_query: 字符串形式的sql查询语句，用于提取mysql中的某张表。
    :return: 表格读取和保存结果
    """
    print("call extract data")
    connection = pymysql.connect(host=host, user=user, password=password, db=db, port=int(port), charset="utf8")

    try:
        # 执行sql并保存为全局变量
        df = pd.read_sql(sql_query, connection)
        # globals()[df_name] = df
        print(df)
        return f"成功创建pandas对象df，包含从mysql中提取的数据。"
    except Exception as e:
        return f"提取数据失败：{e}"
    finally:
        connection.close()


# python代码执行工具
class PythonCodeInput(BaseModel):
    py_code: str = Field(description="一段合法的python代码字符串，用于执行。")


@tool(args_schema=PythonCodeInput)
def python_inter(py_code: str) -> str:
    """
    当用户需要执行python代码时，比如对data frame进行数据分析。请使用此工具。
    ：param py_code：要执行的python代码的字符串
    ：return：返回这段python代码的运行结果
    """
    return str(eval(py_code))


graph = create_react_agent(
    model=llm,
    tools=[extract_data, sql_inter],
    prompt='你是一个智能AI助手，名字是小团团。你的性格活泼、开朗、乐观，是大家的开心果。请用符合你性格的语言回答问题。如果有不知道的问题，可以调用相关的tool辅助搜索答案。'
)
